#include "IOCTLSerializer.h"

ULONG IOCTLSerializer::SerializeEmptyPacket(UCHAR type, BOOLEAN isCompleted, PIRP pIrp)
{
	ULONG bytesWrote = 0;
	bytesWrote += buffer.Write(type);
	bytesWrote += buffer.Write(isCompleted);
	if (isCompleted)
	{
		bytesWrote += buffer.Write(pIrp->IoStatus.Status);
	}
	return bytesWrote;
}

ULONG IOCTLSerializer::SerializeGUID(GUID *pGUID)
{
	ULONG bytesWrote = 0;
	bytesWrote += buffer.Write(&pGUID->Data1, sizeof(pGUID->Data1));
	bytesWrote += buffer.Write(&pGUID->Data2, sizeof(pGUID->Data2));
	bytesWrote += buffer.Write(&pGUID->Data3, sizeof(pGUID->Data3));
	bytesWrote += buffer.Write(pGUID->Data4,  sizeof(pGUID->Data4));
	return bytesWrote;
}

#if (_WIN32_WINNT >= 0x0600)
ULONG IOCTLSerializer::SerializeUsbIdString(PUSB_ID_STRING pString)
{
	ULONG bytesWrote = 0;
	bytesWrote += buffer.Write(&pString->LengthInBytes, sizeof(pString->LengthInBytes));
	bytesWrote += buffer.Write(pString->Buffer, pString->LengthInBytes);
	return bytesWrote;
}
#endif

ULONG IOCTLSerializer::SerializeIOCTL(PIRP pIrp, PVOID pBuffer, ULONG bufferSize, BOOLEAN isCompleted)
{
	ULONG i;
	ULONG bytesWrote = 0;
	buffer.SetBuffer(pBuffer, bufferSize);
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(pIrp);
	ULONG ioctl = stack->Parameters.DeviceIoControl.IoControlCode;

	switch (ioctl)
	{
	case IOCTL_INTERNAL_USB_CYCLE_PORT:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_INTERNAL_USB_CYCLE_PORT, isCompleted, pIrp);
		}
		break;

	case IOCTL_INTERNAL_USB_GET_BUS_INFO:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_INTERNAL_USB_GET_BUS_INFO, isCompleted, pIrp);

			PUSB_BUS_NOTIFICATION pBusNotif = (PUSB_BUS_NOTIFICATION) stack->Parameters.Others.Argument1;
			if ((pBusNotif != NULL) && isCompleted)
			{
				bytesWrote += buffer.Write(&pBusNotif->TotalBandwidth,       sizeof(pBusNotif->TotalBandwidth));
				bytesWrote += buffer.Write(&pBusNotif->ConsumedBandwidth,    sizeof(pBusNotif->ConsumedBandwidth));
				bytesWrote += buffer.Write(&pBusNotif->ControllerNameLength, sizeof(pBusNotif->ControllerNameLength));
			}
		}
		break;

	case IOCTL_INTERNAL_USB_GET_CONTROLLER_NAME:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_INTERNAL_USB_GET_CONTROLLER_NAME, isCompleted, pIrp);

			PUSB_HUB_NAME pHubName = (PUSB_HUB_NAME) stack->Parameters.Others.Argument1;
			if ((pHubName != NULL) && isCompleted)
			{
				bytesWrote += buffer.Write(&pHubName->ActualLength, sizeof(pHubName->ActualLength));
				for (i=0; i<(pHubName->ActualLength-sizeof(ULONG))/2; ++i)
				{
					bytesWrote += buffer.Write(&pHubName->HubName[i],  sizeof(pHubName->HubName[i]));
				}
			}
		}
		break;

	case IOCTL_INTERNAL_USB_GET_HUB_NAME:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_INTERNAL_USB_GET_HUB_NAME, isCompleted, pIrp);

			PUSB_HUB_NAME pHubName = (PUSB_HUB_NAME) pIrp->AssociatedIrp.SystemBuffer;
			if ((pHubName != NULL) && isCompleted)
			{
				bytesWrote += buffer.Write(&pHubName->ActualLength, sizeof(pHubName->ActualLength));
				for (i=0; i<(pHubName->ActualLength-sizeof(ULONG))/2; ++i)
				{
					bytesWrote += buffer.Write(&pHubName->HubName[i], sizeof(pHubName->HubName[i]));
				}
			}
		}
		break;

	case IOCTL_INTERNAL_USB_GET_PORT_STATUS:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_INTERNAL_USB_GET_PORT_STATUS, isCompleted, pIrp);
			if (isCompleted)
			{
				bytesWrote += buffer.Write(stack->Parameters.Others.Argument1, sizeof(ULONG));
			}
		}
		break;

	case IOCTL_INTERNAL_USB_RESET_PORT:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_INTERNAL_USB_RESET_PORT, isCompleted, pIrp);
		}
		break;

	case IOCTL_INTERNAL_USB_SUBMIT_IDLE_NOTIFICATION:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_INTERNAL_USB_SUBMIT_IDLE_NOTIFICATION, isCompleted, pIrp);
		}
		break;

	case IOCTL_USB_GET_DESCRIPTOR_FROM_NODE_CONNECTION:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_USB_GET_DESCRIPTOR_FROM_NODE_CONNECTION, isCompleted, pIrp);
		}
		break;
#if (_WIN32_WINNT < 0x0600)
	case IOCTL_USB_HUB_CYCLE_PORT: // not supported on Windows Vista and later
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_USB_HUB_CYCLE_PORT, isCompleted, pIrp);
		}
		break;

	case IOCTL_USB_GET_HUB_CAPABILITIES: // replaced with IOCTL_USB_GET_HUB_CAPABILITIES_EX on Vista
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_USB_GET_HUB_CAPABILITIES, isCompleted, pIrp);
			PUSB_HUB_CAPABILITIES pHubCap = (PUSB_HUB_CAPABILITIES) pIrp->AssociatedIrp.SystemBuffer;
			if ((pHubCap != NULL) && isCompleted)
			{
				BOOLEAN HubIs2xCapable = (BOOLEAN) pHubCap->HubIs2xCapable;
				bytesWrote += buffer.Write(HubIs2xCapable);
			}			
		}
		break;
#else
	case IOCTL_USB_GET_HUB_CAPABILITIES_EX:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_USB_GET_HUB_CAPABILITIES_EX, isCompleted, pIrp);

			PUSB_HUB_CAPABILITIES_EX pHubCapEx = (PUSB_HUB_CAPABILITIES_EX) pIrp->AssociatedIrp.SystemBuffer;
			if ((pHubCapEx != NULL) && isCompleted)
			{
				bytesWrote += buffer.Write(&pHubCapEx->CapabilityFlags.ul, sizeof(ULONG));
			}
		}
		break;

	case IOCTL_INTERNAL_USB_GET_DEVICE_CONFIG_INFO:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_INTERNAL_USB_GET_DEVICE_CONFIG_INFO, isCompleted, pIrp);

			PHUB_DEVICE_CONFIG_INFO pHDCI = (PHUB_DEVICE_CONFIG_INFO) stack->Parameters.Others.Argument1;
			if ((pHDCI != NULL) && isCompleted)
			{
				bytesWrote += buffer.Write(&pHDCI->Version,     sizeof(pHDCI->Version));
				bytesWrote += buffer.Write(&pHDCI->Length,      sizeof(pHDCI->Length));
				bytesWrote += buffer.Write(&pHDCI->HubFlags.ul, sizeof(pHDCI->HubFlags.ul));

				bytesWrote += SerializeUsbIdString(&pHDCI->HardwareIds);
				bytesWrote += SerializeUsbIdString(&pHDCI->CompatibleIds);
				bytesWrote += SerializeUsbIdString(&pHDCI->DeviceDescription);

				PUSB_HUB_DEVICE_UXD_SETTINGS pUxd = &pHDCI->UxdSettings;
				bytesWrote += buffer.Write(&pUxd->Version, sizeof(pUxd->Version));
				bytesWrote += SerializeGUID(&pUxd->PnpGuid);
				bytesWrote += SerializeGUID(&pUxd->OwnerGuid);
				bytesWrote += buffer.Write(&pUxd->DeleteOnShutdown,   sizeof(pUxd->DeleteOnShutdown));
				bytesWrote += buffer.Write(&pUxd->DeleteOnReload,     sizeof(pUxd->DeleteOnReload));
				bytesWrote += buffer.Write(&pUxd->DeleteOnDisconnect, sizeof(pUxd->DeleteOnDisconnect));
			}
		}
		break;

	case IOCTL_INTERNAL_USB_GET_TOPOLOGY_ADDRESS:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_INTERNAL_USB_GET_TOPOLOGY_ADDRESS, isCompleted, pIrp);

			PUSB_TOPOLOGY_ADDRESS pTorAddr = (PUSB_TOPOLOGY_ADDRESS) stack->Parameters.Others.Argument1;
			if ((pTorAddr != NULL) && isCompleted)
			{
				bytesWrote += buffer.Write(&pTorAddr->PciBusNumber,      sizeof(pTorAddr->PciBusNumber));
				bytesWrote += buffer.Write(&pTorAddr->PciDeviceNumber,   sizeof(pTorAddr->PciDeviceNumber));
				bytesWrote += buffer.Write(&pTorAddr->PciFunctionNumber, sizeof(pTorAddr->PciFunctionNumber));
				bytesWrote += buffer.Write(&pTorAddr->RootHubPortNumber, sizeof(pTorAddr->RootHubPortNumber));
				bytesWrote += buffer.Write( pTorAddr->HubPortNumber,     sizeof(pTorAddr->HubPortNumber));
			}
		}
		break;
#endif
	case IOCTL_USB_GET_NODE_CONNECTION_INFORMATION:
	case IOCTL_USB_GET_NODE_CONNECTION_INFORMATION_EX:
		{
			if (ioctl == PACK_IOCTL_USB_GET_NODE_CONNECTION_INFORMATION)
			{
				bytesWrote += SerializeEmptyPacket((UCHAR)
					   PACK_IOCTL_USB_GET_NODE_CONNECTION_INFORMATION, isCompleted, pIrp);
			}
			else
			{
				bytesWrote += SerializeEmptyPacket((UCHAR)
					PACK_IOCTL_USB_GET_NODE_CONNECTION_INFORMATION_EX, isCompleted, pIrp);
			}

			PUSB_NODE_CONNECTION_INFORMATION pNCInf = (PUSB_NODE_CONNECTION_INFORMATION)
														pIrp->AssociatedIrp.SystemBuffer;
			bytesWrote += buffer.Write(&pNCInf->ConnectionIndex, sizeof(pNCInf->ConnectionIndex));
			if (!isCompleted)
			{				
				break;
			}

			PUSB_DEVICE_DESCRIPTOR pDevDesc = (PUSB_DEVICE_DESCRIPTOR) &pNCInf->DeviceDescriptor;
			bytesWrote += buffer.Write(&pDevDesc->bLength,            sizeof(pDevDesc->bLength));
			bytesWrote += buffer.Write(&pDevDesc->bDescriptorType,    sizeof(pDevDesc->bDescriptorType));
			bytesWrote += buffer.Write(&pDevDesc->bcdUSB,             sizeof(pDevDesc->bcdUSB));
			bytesWrote += buffer.Write(&pDevDesc->bDeviceClass,       sizeof(pDevDesc->bDeviceClass));
			bytesWrote += buffer.Write(&pDevDesc->bDeviceSubClass,    sizeof(pDevDesc->bDeviceSubClass));
			bytesWrote += buffer.Write(&pDevDesc->bDeviceProtocol,    sizeof(pDevDesc->bDeviceProtocol));
			bytesWrote += buffer.Write(&pDevDesc->bMaxPacketSize0,    sizeof(pDevDesc->bMaxPacketSize0));
			bytesWrote += buffer.Write(&pDevDesc->idVendor,           sizeof(pDevDesc->idVendor));
			bytesWrote += buffer.Write(&pDevDesc->idProduct,          sizeof(pDevDesc->idProduct));
			bytesWrote += buffer.Write(&pDevDesc->bcdDevice,          sizeof(pDevDesc->bcdDevice));
			bytesWrote += buffer.Write(&pDevDesc->iManufacturer,      sizeof(pDevDesc->iManufacturer));
			bytesWrote += buffer.Write(&pDevDesc->iProduct,           sizeof(pDevDesc->iProduct));
			bytesWrote += buffer.Write(&pDevDesc->iSerialNumber,      sizeof(pDevDesc->iSerialNumber));
			bytesWrote += buffer.Write(&pDevDesc->bNumConfigurations, sizeof(pDevDesc->bNumConfigurations));

			bytesWrote += buffer.Write(&pNCInf->CurrentConfigurationValue, sizeof(pNCInf->CurrentConfigurationValue));
			bytesWrote += buffer.Write(&pNCInf->LowSpeed,                  sizeof(pNCInf->LowSpeed));
			bytesWrote += buffer.Write(&pNCInf->DeviceIsHub,               sizeof(pNCInf->DeviceIsHub));
			bytesWrote += buffer.Write(&pNCInf->DeviceAddress,             sizeof(pNCInf->DeviceAddress));
			bytesWrote += buffer.Write(&pNCInf->NumberOfOpenPipes,         sizeof(pNCInf->NumberOfOpenPipes));
			bytesWrote += buffer.Write(&pNCInf->ConnectionStatus,          sizeof(pNCInf->ConnectionStatus));

			for (i=0; i<pNCInf->NumberOfOpenPipes; ++i)
			{
				PUSB_ENDPOINT_DESCRIPTOR pED = (PUSB_ENDPOINT_DESCRIPTOR) &pNCInf->PipeList[i].EndpointDescriptor;
				bytesWrote += buffer.Write(&pED->bLength,          sizeof(pED->bLength));
				bytesWrote += buffer.Write(&pED->bDescriptorType,  sizeof(pED->bDescriptorType));
				bytesWrote += buffer.Write(&pED->bEndpointAddress, sizeof(pED->bEndpointAddress));
				bytesWrote += buffer.Write(&pED->bmAttributes,     sizeof(pED->bmAttributes));
				bytesWrote += buffer.Write(&pED->wMaxPacketSize,   sizeof(pED->wMaxPacketSize));
				bytesWrote += buffer.Write(&pED->bInterval,        sizeof(pED->bInterval));

				bytesWrote += buffer.Write(&pNCInf->PipeList[i].ScheduleOffset, sizeof(ULONG));
			}
		}
		break;

	case IOCTL_USB_GET_NODE_CONNECTION_NAME:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_USB_GET_NODE_CONNECTION_NAME, isCompleted, pIrp);

			PUSB_NODE_CONNECTION_NAME pNodeName = (PUSB_NODE_CONNECTION_NAME) pIrp->AssociatedIrp.SystemBuffer;
			if (pNodeName == NULL)
			{
				break;
			}
			bytesWrote += buffer.Write(&pNodeName->ConnectionIndex, sizeof(pNodeName->ConnectionIndex));
			if (!isCompleted)
			{
				break;
			}				
			bytesWrote += buffer.Write(&pNodeName->ActualLength,    sizeof(pNodeName->ActualLength));
			for (i=0; i<pNodeName->ActualLength/2; ++i)
			{
				bytesWrote += buffer.Write(&pNodeName->NodeName[i], sizeof(pNodeName->NodeName[i]));
			}
		}
		break;

	case IOCTL_USB_GET_NODE_INFORMATION:
		{
			bytesWrote += SerializeEmptyPacket((UCHAR)PACK_IOCTL_USB_GET_NODE_INFORMATION, isCompleted, pIrp);
			PUSB_NODE_INFORMATION pNodeInf = (PUSB_NODE_INFORMATION) pIrp->AssociatedIrp.SystemBuffer;
			if ((pNodeInf != NULL) && isCompleted)
			{
				bytesWrote += buffer.Write(&pNodeInf->NodeType, sizeof(pNodeInf->NodeType));
				if (pNodeInf->NodeType == UsbMIParent)
				{
					PUSB_MI_PARENT_INFORMATION pMiInf = (PUSB_MI_PARENT_INFORMATION) &pNodeInf->u;
					bytesWrote += buffer.Write(&pMiInf->NumberOfInterfaces, sizeof(ULONG));
					break;
				}

				PUSB_HUB_DESCRIPTOR pHubDesc = &((PUSB_HUB_INFORMATION) &pNodeInf->u)->HubDescriptor;

				bytesWrote += buffer.Write(&pHubDesc->bDescriptorLength,   sizeof(pHubDesc->bDescriptorLength));
				bytesWrote += buffer.Write(&pHubDesc->bDescriptorType,     sizeof(pHubDesc->bDescriptorType));
				bytesWrote += buffer.Write(&pHubDesc->bNumberOfPorts,      sizeof(pHubDesc->bNumberOfPorts));
				bytesWrote += buffer.Write(&pHubDesc->wHubCharacteristics, sizeof(pHubDesc->wHubCharacteristics));
				bytesWrote += buffer.Write(&pHubDesc->bPowerOnToPowerGood, sizeof(pHubDesc->bPowerOnToPowerGood));
				bytesWrote += buffer.Write(&pHubDesc->bHubControlCurrent,  sizeof(pHubDesc->bHubControlCurrent));

				bytesWrote += buffer.Write(((PUSB_HUB_INFORMATION) &pNodeInf->u)->HubIsBusPowered);				
			}
		}
		break;

	case IOCTL_INTERNAL_USB_SUBMIT_URB:
		{
			PURB pUrb = (PURB) stack->Parameters.Others.Argument1;
			bytesWrote = urbSerializer.SerializeURB(pUrb, pBuffer, bufferSize, isCompleted);
		}
		break;

	default:
		KdPrint(("IOCTLSerializer::SerializeIOCTL(...): Unknown IO Control Code 0x%x\n", ioctl));
		break;
	}
	return bytesWrote;
}